#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用户到达地址表生成器

基于addr_info.csv和家庭宽带办理信息表.csv，按照工作日和休息日的行为模式
为每个家庭生成用户轨迹数据。

工作日模式：住宅->(公司/学校)->(住宅/餐厅)->(公司/学校)->(超市)->住宅
休息日模式：住宅->(景点/影院)->(住宅/餐厅)->(景点/影院/4S店)->(超市)->住宅

输出格式：user_id, time, arrive_addr, add_type, longitude, latitude, stay_time, lac_code
"""

import csv
import json
import random
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional
from config import dummy_data_dir


class AddressManager:
    """地址数据管理器"""
    
    def __init__(self, addr_file: str):
        self.addr_data = {}
        self.load_addr_data(addr_file)
        
    def load_addr_data(self, addr_file: str):
        """加载地址信息"""
        try:
            with open(addr_file, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    name = row['name']
                    self.addr_data[name] = {
                        'type': row['type'],
                        'longitude': float(row['longitude']),
                        'latitude': float(row['latitude']),
                        'lac_code': row['lac_code']
                    }
        except Exception as e:
            print(f"加载地址数据失败: {e}")
            raise
            
    def get_addr_info(self, name: str) -> Optional[Dict]:
        """获取地址详细信息"""
        return self.addr_data.get(name)


class FamilyManager:
    """家庭信息管理器"""
    
    def __init__(self, broadband_file: str):
        self.families = {}
        self.load_broadband_data(broadband_file)
        
    def load_broadband_data(self, broadband_file: str):
        """加载宽带信息，按地址聚合家庭"""
        try:
            with open(broadband_file, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    user_id = row['user_id']
                    brd_addr = row['brd_addr']
                    
                    # 提取小区名称
                    parts = brd_addr.split('栋')
                    if len(parts) >= 2:
                        community = parts[0].replace('区', '')
                        
                        if community not in self.families:
                            self.families[community] = []
                        self.families[community].append({
                            'user_id': user_id,
                            'full_addr': brd_addr,
                            'community': community
                        })
        except Exception as e:
            print(f"加载宽带数据失败: {e}")
            raise
            
    def get_families(self) -> Dict[str, List[Dict]]:
        """获取所有家庭信息"""
        return self.families


class UserTrajectoryGenerator:
    """用户轨迹生成器"""
    
    def __init__(self, addr_manager: AddressManager):
        self.addr_manager = addr_manager
        
    def get_available_addresses(self, addr_type: str) -> List[Dict]:
        """获取指定类型的可用地址"""
        addresses = []
        for name, info in self.addr_manager.addr_data.items():
            if info['type'] == addr_type:
                addresses.append({
                    'name': name,
                    'type': addr_type,
                    'longitude': info['longitude'],
                    'latitude': info['latitude'],
                    'lac_code': info['lac_code']
                })
        return addresses
        
    def generate_stay_time(self, addr_type: str) -> int:
        """根据地址类型生成合理的驻留时长（分钟）"""
        stay_times = {
            '住宅': random.randint(360, 720),  # 6-12小时
            '公司': random.randint(360, 480),  # 6-8小时
            '小学': random.randint(240, 360),  # 4-6小时
            '中学': random.randint(300, 420),  # 5-7小时
            '大学': random.randint(180, 480),  # 3-8小时
            '餐厅': random.randint(30, 120),   # 30分钟-2小时
            '景点': random.randint(120, 300),  # 2-5小时
            '超市': random.randint(20, 60),    # 20-60分钟
            '影院': random.randint(90, 180),   # 1.5-3小时
            '4S店': random.randint(60, 180)    # 1-3小时
        }
        return stay_times.get(addr_type, random.randint(30, 120))
        
    def generate_time_sequence(self, pattern: List[str], date: datetime) -> List[datetime]:
        """根据行为模式生成时间序列"""
        times = []
        current_time = datetime(date.year, date.month, date.day, 6, 0)  # 从早上6点开始
        
        for i, addr_type in enumerate(pattern):
            if i == 0:  # 第一个点（住宅）
                times.append(current_time)
            elif addr_type == '住宅' and i == len(pattern) - 1:  # 最后一个住宅
                # 晚上回家时间
                current_time = datetime(date.year, date.month, date.day, 19, 0) + \
                             timedelta(minutes=random.randint(-60, 120))
                times.append(current_time)
            elif addr_type in ['公司', '小学', '中学', '大学']:  # 工作/学习时间
                if i == 1:  # 上午
                    current_time = datetime(date.year, date.month, date.day, 8, 0) + \
                                 timedelta(minutes=random.randint(-30, 60))
                else:  # 下午
                    current_time = datetime(date.year, date.month, date.day, 13, 0) + \
                                 timedelta(minutes=random.randint(0, 120))
                times.append(current_time)
            elif addr_type == '餐厅':  # 用餐时间
                if i == 2:  # 午餐
                    current_time = datetime(date.year, date.month, date.day, 12, 0) + \
                                 timedelta(minutes=random.randint(-30, 30))
                else:  # 晚餐
                    current_time = datetime(date.year, date.month, date.day, 18, 0) + \
                                 timedelta(minutes=random.randint(-30, 30))
                times.append(current_time)
            elif addr_type == '超市':  # 购物时间
                current_time = datetime(date.year, date.month, date.day, 17, 0) + \
                             timedelta(minutes=random.randint(-60, 120))
                times.append(current_time)
            elif addr_type in ['景点', '影院', '4S店']:  # 休闲娱乐时间
                if i == 1:  # 上午
                    current_time = datetime(date.year, date.month, date.day, 9, 0) + \
                                 timedelta(minutes=random.randint(0, 120))
                else:  # 下午
                    current_time = datetime(date.year, date.month, date.day, 14, 0) + \
                                 timedelta(minutes=random.randint(0, 180))
                times.append(current_time)
                
        return times
        
    def generate_user_trajectory(self, user_id: str, full_home_addr: str, start_date: datetime) -> List[Dict]:
        """为单个用户生成7天轨迹"""
        trajectory = []
        
        # 获取住宅地址信息 - 直接使用宽带表中的完整地址
        home_info = None
        
        # 从addr_info.csv中查找对应的住宅地址信息
        for name, info in self.addr_manager.addr_data.items():
            if info['type'] == '住宅' and name in full_home_addr:
                home_info = {
                    'name': full_home_addr,  # 使用完整的住宅地址
                    'type': '住宅',
                    'longitude': info['longitude'],
                    'latitude': info['latitude'],
                    'lac_code': info['lac_code']
                }
                break
                
        if not home_info:
            # 如果没有精确匹配，使用小区名称匹配
            community = full_home_addr.split('栋')[0] if '栋' in full_home_addr else full_home_addr
            home_addr = self.get_available_addresses('住宅')
            for addr in home_addr:
                if community in addr['name']:
                    home_info = {
                        'name': full_home_addr,  # 使用完整的住宅地址
                        'type': '住宅',
                        'longitude': addr['longitude'],
                        'latitude': addr['latitude'],
                        'lac_code': addr['lac_code']
                    }
                    break
                    
        if not home_info:
            # 如果都没有匹配，随机选择一个住宅地址
            home_addr = self.get_available_addresses('住宅')
            if home_addr:
                addr = random.choice(home_addr)
                home_info = {
                    'name': full_home_addr,  # 使用完整的住宅地址
                    'type': '住宅',
                    'longitude': addr['longitude'],
                    'latitude': addr['latitude'],
                    'lac_code': addr['lac_code']
                }
            else:
                return trajectory
            
        # 为用户分配固定的角色类型
        user_profile = self.assign_user_profile()
        
        # 根据用户角色类型生成对应的轨迹模式
        def get_user_pattern(profile: Dict, is_weekday: bool) -> List[str]:
            """根据用户角色和日期类型生成轨迹模式"""
            addr_type = profile['addr_type']
            
            if is_weekday:
                # 工作日模式
                if addr_type == '公司':
                    patterns = [
                        ['住宅', '公司', '住宅', '公司', '超市', '住宅'],
                        ['住宅', '公司', '餐厅', '公司', '超市', '住宅']
                    ]
                elif addr_type in ['小学', '中学', '大学']:
                    patterns = [
                        ['住宅', addr_type, '住宅', addr_type, '超市', '住宅'],
                        ['住宅', addr_type, '餐厅', addr_type, '超市', '住宅']
                    ]
                else:
                    patterns = [['住宅', '公司', '住宅', '公司', '超市', '住宅']]
            else:
                # 休息日模式（所有角色通用）
                patterns = [
                    ['住宅', '景点', '住宅', '景点', '超市', '住宅'],
                    ['住宅', '影院', '住宅', '影院', '超市', '住宅'],
                    ['住宅', '景点', '餐厅', '景点', '超市', '住宅'],
                    ['住宅', '景点', '餐厅', '4S店', '超市', '住宅']
                ]
            
            return random.choice(patterns)
        
        # 生成7天轨迹
        for day in range(7):
            current_date = start_date + timedelta(days=day)
            
            # 判断是工作日还是休息日
            is_weekday = current_date.weekday() < 5  # 周一到周五
            pattern = get_user_pattern(user_profile, is_weekday)
                
            # 生成时间序列
            times = self.generate_time_sequence(pattern, current_date)
            
            # 生成轨迹记录
            for i, (addr_type, time) in enumerate(zip(pattern, times)):
                if addr_type == '住宅':
                    addr_info = home_info
                else:
                    # 根据用户档案选择对应的地址
                    available = self.get_available_addresses(addr_type)
                    if available:
                        if addr_type in ['公司', '小学', '中学', '大学']:
                            # 使用用户固定的工作/学校地址
                            addr_info = self.select_fixed_address(user_profile, addr_type, available)
                        else:
                            addr_info = random.choice(available)
                    else:
                        continue
                        
                record = {
                    'user_id': user_id,
                    'time': time.strftime('%Y-%m-%d %H:%M:%S'),
                    'arrive_addr': addr_info['name'],
                    'add_type': addr_type,
                    'longitude': addr_info['longitude'],
                    'latitude': addr_info['latitude'],
                    'stay_time': self.generate_stay_time(addr_type),
                    'lac_code': addr_info['lac_code']
                }
                trajectory.append(record)
                
        return trajectory
        
    def assign_user_profile(self) -> Dict:
        """为用户分配固定的角色类型和对应地址"""
        # 定义用户角色类型及其权重
        role_types = [
            {'type': '上班族', 'addr_type': '公司', 'weight': 0.6},
            {'type': '小学生', 'addr_type': '小学', 'weight': 0.15},
            {'type': '中学生', 'addr_type': '中学', 'weight': 0.15},
            {'type': '大学生', 'addr_type': '大学', 'weight': 0.1}
        ]
        
        # 根据权重随机选择角色
        total_weight = sum(role['weight'] for role in role_types)
        rand_val = random.uniform(0, total_weight)
        cumulative = 0
        selected_role = None
        
        for role in role_types:
            cumulative += role['weight']
            if rand_val <= cumulative:
                selected_role = role
                break
        
        # 获取对应类型的地址
        addr_type = selected_role['addr_type']
        available = self.get_available_addresses(addr_type)
        
        if available:
            fixed_address = random.choice(available)
        else:
            # 如果没有对应类型的地址，使用公司类型
            company_addrs = self.get_available_addresses('公司')
            fixed_address = random.choice(company_addrs) if company_addrs else {}
            addr_type = '公司'
            
        return {
            'role': selected_role['type'],
            'addr_type': addr_type,
            'fixed_address': fixed_address
        }
        
    def select_fixed_address(self, profile: Dict, addr_type: str, available: List[Dict]) -> Dict:
        """选择用户的固定地址"""
        if addr_type == profile['addr_type']:
            return profile['fixed_address']
        return random.choice(available) if available else {}


def main():
    """主函数"""
    # 文件路径
    addr_file = os.path.join(dummy_data_dir, 'addr_info.csv').replace(os.sep, '/')
    broadband_file = os.path.join(dummy_data_dir, '家庭宽带办理信息表.csv').replace(os.sep, '/')
    output_file = os.path.join(dummy_data_dir, '用户到达地址表.csv').replace(os.sep, '/')
    
    try:
        # 初始化管理器
        addr_manager = AddressManager(addr_file)
        family_manager = FamilyManager(broadband_file)
        generator = UserTrajectoryGenerator(addr_manager)
        
        # 获取所有家庭
        families = family_manager.get_families()
        
        # 生成轨迹数据
        all_trajectories = []
        start_date = datetime(2024, 1, 15)  # 从2024年1月15日开始
        
        for community, members in families.items():
            print(f"正在处理小区: {community}, 用户数: {len(members)}")
            
            for member in members:
                user_id = member['user_id']
                full_home_addr = member['full_addr']  # 使用完整的住宅地址
                trajectory = generator.generate_user_trajectory(
                    user_id, full_home_addr, start_date
                )
                all_trajectories.extend(trajectory)
                
        # 写入CSV文件
        with open(output_file, 'w', encoding='utf-8', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=[
                'user_id', 'time', 'arrive_addr', 'add_type', 
                'longitude', 'latitude', 'stay_time', 'lac_code'
            ])
            writer.writeheader()
            writer.writerows(all_trajectories)
            
        print(f"轨迹生成完成！共生成 {len(all_trajectories)} 条记录")
        print(f"输出文件: {output_file}")
        
    except Exception as e:
        print(f"生成轨迹时发生错误: {e}")
        raise


if __name__ == "__main__":
    main()